/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2024. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#ifndef GE_TILING_CONTEXT_H
#define GE_TILING_CONTEXT_H
#include "graph/tensor.h"
#include "graph/graph_api_export.h"
#include "register/tiling_data.h"
#include "register/storage_shape.h"
#include "register/runtime_context.h"
#include "register/definition.h"

namespace fe {
class PlatFormInfos;
}  // namespace ge

namespace ge {
/**
 * tiling kernel的context
 */
class GRAPH_API_EXPORT TilingContext : public RuntimeContext {
public:
    TilingContext(ge::Node& node);
    const void *GetCompileInfo() const;
    bool HasError() const
    {
        return hasError_;
    }
    /**
     * 获取CompileInfo
     * @tparam T CompileInfo的类型
     * @return CompileInfo的指针
     */
    template <typename T>
    const T *GetCompileInfo() const
    {
        return reinterpret_cast<const T *>(GetCompileInfo());
    }
    /**
     * 获取 fe::PlatFormInfos 指针
     * @return fe::PlatFormInfos 指针
     */
    fe::PlatFormInfos *GetPlatformInfo() const;
    /**
     * 设置tiling key
     * @param tiling_key tiling key
     * @return 成功时返回ge::GRAPH_SUCCESS
     */
    ge::graphStatus SetTilingKey(const uint64_t tiling_key);
    /**
     * 获取tiling key
     * @return tiling key，获取失败时返回UINT64_MAX
     */
    uint64_t GetTilingKey() const;
    /**
     * 设置block dim
     * @param block_dim block dim
     * @return 成功时返回ge::GRAPH_SUCCESS
     */
    ge::graphStatus SetBlockDim(const uint32_t block_dim);
    /**
     * 获取block dim
     * @return block dim
     */
    uint32_t GetBlockDim() const;
    /**
     * 获取输入shape，输入shape中包含了原始shape与运行时shape
     * @param index 输入index
     * @return 输入shape指针，index非法时返回空指针
     */
    const StorageShape *GetInputShape(const size_t index) const;
    /**
     * 设置schedule_mode
     * @param schedule_mode schedule_mode
     * @return 成功时返回ge::GRAPH_SUCCESS
     */
    ge::graphStatus SetScheduleMode(const uint32_t schedule_mode);
    /**
     * 获取设置schedule_mode
     * @return 设置schedule_mode，获取失败时
     */
    uint32_t GetScheduleMode() const;
    /**
     * 获取workspace sizes指针
     * @param workspace_count workspace的个数，传入的workspace个数不可以超过编译时指定的最大workspace个数
     * @return workspace sizes指针
     */
    size_t *GetWorkspaceSizes(const size_t workspace_count);
    /**
     * 获取 workspace 个数
     * @return workspace 个数
     */
    size_t GetWorkspaceNum() const;
    /**
     * 获取有类型的tiling data指针
     * @tparam T tiling data类型，sizeof(T)不可以大于编译结果中指定的最大tiling data长度
     * @return tiling data指针，失败时返回空指针
     */
    template <typename T>
    auto GetTilingData() -> T *
    {
        auto tiling_data = GetRawTilingData();
        if (tiling_data == nullptr) {
            return nullptr;
        }
        if (tiling_data->GetCapacity() < sizeof(T)) {
            return nullptr;
        }
        tiling_data->SetDataSize(sizeof(T));
        return static_cast<T *>(tiling_data->GetData());
    }
    /**
     * 获取无类型的tiling data码流
     * @return tiling data指针，失败时返回空指针
     */
    TilingData *GetRawTilingData();
    /**
     * 根据输出index，获取输出shape指针，shape中包含了原始shape与运行时shape
     * @param index 输出index
     * @return 输出shape指针，index非法时，返回空指针
     */
    const StorageShape *GetOutputShape(const size_t index) const;
    /**
     * 基于算子IR原型定义，获取`OPTIONAL_INPUT`类型的输入shape指针，shape中包含了原始shape与运行时shape
     * WARN：受限于代码架构，对于内置算子，该函数不支持，返回结果不可信
     * @param ir_index IR原型定义中的index
     * @return shape指针，index非法，或该INPUT没有实例化时，返回空指针
     */
    const StorageShape *GetOptionalInputShape(const size_t ir_index) const;
    /**
     * 基于算子IR原型定义，获取`DYNAMIC_INPUT`类型的输入shape指针，shape中包含了原始shape与运行时shape
     * WARN：受限于代码架构，对于内置算子，该函数不支持
     * @param ir_index IR原型定义中的index
     * @param relative_index
     * 该输入实例化后的相对index，例如某个DYNAMIC_INPUT实例化了3个输入，那么relative_index的有效范围是[0,2]
     * @return shape指针，index或relative_index非法时，返回结果不可信
     */
    const StorageShape *GetDynamicInputShape(const size_t ir_index, const size_t relative_index) const;

    /**
     * 设置tiling cond
     * @param tiling_cond tiling condition
     * @return 成功时返回ge::GRAPH_SUCCESS
     */
    ge::graphStatus SetTilingCond(int32_t tiling_cond);
    /**
     * 获取tiling cond
     * @return tiling cond:有效的tiling_cond大于等于0，若该值无效返回-1
     */
    int32_t GetTilingCond() const;
    /**
     * 设置是否需要atomic clean
     * @param atomic true/false代表是否需要做atomic clean
     * @return 成功时返回ge::GRAPH_SUCCESS
     */
    ge::graphStatus SetNeedAtomic(const bool atomic);
    /**
     * 获取是否需要atomic clean
     * @return true/false
     */
    bool NeedAtomic() const;

    /**
     * 获取 确定性计算变量
     * @return int32 变量
     */
    int32_t GetDeterministic() const;
    RuntimeAttrs *GetAttrs() override;
    // 设置当前node是否为自定义算子
    TilingContext &IsCustomOp(bool isCustomOp);
    // 设置当前算子的所有IR属性名称
    TilingContext &OriginAttrNames(const std::vector<std::string> &attrNames);

    /**
     * 设置 fe::PlatFormInfos 指针
     * @return void
     */
    void SetPlatformInfo(fe::PlatFormInfos *platform_info);

protected:
    size_t GetRealIndex(const size_t ir_index, const size_t relative_index) const override;

private:
    void SaveRuntimeAttrs();
    RuntimeAttrs *GetAttrsInternal();
    static constexpr size_t MAX_WORKSPACE_NUMBER = 128;

    bool isCustomOp_ = false;
    bool hasError_ = false;
    // 实际数据内容是`gert::TilingData`
    std::unique_ptr<uint8_t[]> tilingDataContainer_;
    std::vector<size_t> workspace_;
    size_t workspaceNum_ = 0;
    uint64_t tilingKey_ = 0;
    uint32_t blockDim_ = 0;
    // 默认为normal模式
    uint32_t scheduleMode_ = 0;
    int32_t tilingCond_ = -1;
    bool needAtomic_ = true;
    std::vector<gert::StorageShape> inputShapes_;
    std::vector<gert::StorageShape> outputShapes_;
    std::vector<std::string> originAttrNames_;
    std::unique_ptr<RuntimeAttrs> runtimeAttrs_;
    fe::PlatFormInfos *platformInfos_ = nullptr;
    std::unique_ptr<uint8_t[]> compileInfo_ = nullptr;
    int32_t deterministic_ = INT32_MAX;
};
}  // namespace ge
#endif  // GE_TILING_CONTEXT_H
